import re
import random

def clean_text(text):
    text = re.sub(r'\\boxed{\\text{(.*?)}}', r'\\boxed{\1}', text)
    return text

def clean_text(text):
    """Clean the text by removing unnecessary whitespace and newlines."""
    if "\\text{" in text:
        if text.startswith('\\text{'):
            text = text[6:-1]
    return text

def extract_answer(answer_str):
    if 'boxed' in answer_str:
        ans = answer_str.split('boxed')[-1]
        if len(ans) == 0:
            return ''
        elif ans[0] == '{':
            stack = 1
            a = ''
            for c in ans[1:]:
                if c == '{':
                    stack += 1
                    a += c
                elif c == '}':
                    stack -= 1
                    if stack == 0:
                        break
                    a += c
                else:
                    a += c
            return clean_text(a)
        else:
            return ''
    else:
        return ''


def extract_solution(solution_str):
    # Remove everything before the first "Assistant:"
    # if "Assistant:" in solution_str:
    #     solution_str = solution_str.split("Assistant:", 1)[1]
    # else:
    #     return None

    
    # solution_str= solution_str.split('</think>')[-1].strip()
    # final_answer = extract_answer(solution_str)
    # if final_answer == '':
    #         final_answer = None
    # return final_answer

    answer_pattern = r'<answer>(.*?)</answer>'
    matches = re.findall(answer_pattern, solution_str, re.DOTALL)

    if matches:
        final_answer = matches[-1].strip()
        final_answer = extract_answer(final_answer)
        if final_answer == '':
            final_answer = None
    else:
        final_answer = None


    return final_answer


def compute_score(data_source, solution_str, ground_truth, extra_info, method='strict', format_score=0.1, score=1.):
    """The scoring function for GSM8k.

    Reference: Trung, Luong, et al. "Reft: Reasoning with reinforced fine-tuning." Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2024.

    Args:
        solution_str: the solution text
        ground_truth: the ground truth
        method: the method to extract the solution, choices are 'strict' and 'flexible'
        format_score: the score for the format
        score: the score for the correct answer
    """
    answer = extract_solution(solution_str=solution_str)
    do_print = random.randint(1, 64) == 1
    if do_print:
        print(f"--------------------------------")
        print(f"Ground truth: {ground_truth} | Extracted answer: {answer}")
        print(f"Solution string: {solution_str}")

    if answer is None:
        if do_print:
            print(f"No answer found")
        return 0
    else:
        if answer.strip() == ground_truth.strip():
            if do_print:
                print(f"Correct answer: {answer}")
            return score
        else:
            if do_print:
                print(f"Incorrect answer {answer} | Ground truth: {ground_truth}")
            return format_score
